# %% 导入库
import tensorflow as tf
import matplotlib.pyplot as plt
import os
from tensorflow.keras import layers
import time
from IPython import display
from PIL import Image
import pathlib
import numpy as np

# 设置超参数
EPOCHS = 100  # 训练轮次，减少以加快训练速度
noise_dim = 160  # 噪声向量的维度
num_examples_to_generate = 9  # 每次生成的样本数量
BATCH_SIZE = 512  # 批量大小，增加以提高训练效率

# %% 加载数据 将图片为文件转为数组
dir_path = "anime-faces"  # 数据集目录
data_dir = pathlib.Path(dir_path)  # 使用pathlib加载目录
imgs = list(data_dir.glob('*.png'))  # 获取目录下所有PNG文件
img_arr = []  # 存储图像数组
target_size = (56, 56)  # 图像的目标尺寸

# 遍历所有图片文件，将其转换为灰度图像并调整大小，然后添加到数组中
for img in imgs:
    img = Image.open(str(img)).convert('L')  # 转换为灰度图像
    img = img.resize(target_size, Image.Resampling.LANCZOS)  # 调整图像尺寸
    img_arr.append(np.array(img))

train_images = np.array(img_arr)  # 将列表转换为numpy数组
nums = train_images.shape[0]  # 图像总数

# 输出图像数量和形状
print(f"Total number of images: {nums}")
print(f"Shape of train_images before reshaping: {train_images.shape}")

# 对图像数据进行归一化处理，使其范围在[-1, 1]
train_images = (train_images - 127.5) / 127.5

# 将图像数据转换为TensorFlow张量，并设置为四维（样本数，高度，宽度，颜色通道）
train_images = train_images.reshape(nums, 56, 56, 1).astype('float32')

# 创建TensorFlow数据集，打乱数据并按批次加载
train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(nums).batch(BATCH_SIZE)

# %% 定义生成器模型
def make_generator_model():
    """
    创建生成器模型，该模型从噪声向量生成图像。
    """
    model = tf.keras.Sequential([
        layers.Dense(7 * 7 * 256, use_bias=False, input_shape=(noise_dim,)),  # 输入噪声向量，输出扁平化的特征图
        layers.BatchNormalization(),  # 批量归一化加速训练
        layers.LeakyReLU(),  # 使用LeakyReLU激活函数避免死区问题
        layers.Reshape((7, 7, 256)),  # 重塑为三维特征图
        layers.Conv2DTranspose(128, (5, 5), strides=(2, 2), padding='same', use_bias=False),  # 反卷积操作，扩大图像尺寸
        layers.BatchNormalization(),  # 批量归一化
        layers.LeakyReLU(),  # 激活函数
        layers.Conv2DTranspose(64, (5, 5), strides=(2, 2), padding='same', use_bias=False),  # 再次反卷积
        layers.BatchNormalization(),  # 批量归一化
        layers.LeakyReLU(),  # 激活函数
        layers.Conv2DTranspose(1, (5, 5), strides=(2, 2), padding='same', use_bias=False, activation='tanh')  # 最终输出图像，使用tanh激活函数
    ])
    return model

# 创建生成器实例
generator = make_generator_model()
# 创建一个随机噪声向量用于测试生成器
noise = tf.random.normal([1, noise_dim])

# %% 定义鉴别器模型
def make_discriminator_model():
    """
    创建鉴别器模型，该模型判断输入图像是真实的还是由生成器产生的。
    """
    model = tf.keras.Sequential([
        layers.Conv2D(64, (5, 5), strides=(2, 2), padding='same', input_shape=[56, 56, 1]),  # 输入图像，输出特征图
        layers.LeakyReLU(),  # 激活函数
        layers.Dropout(0.3),  # 防止过拟合
        layers.Conv2D(128, (5, 5), strides=(2, 2), padding='same'),  # 进一步提取特征
        layers.LeakyReLU(),  # 激活函数
        layers.Dropout(0.3),  # 防止过拟合
        layers.Flatten(),  # 展平特征图
        layers.Dense(1)  # 输出一个数值，表示输入图像是真还是假
    ])
    return model

# 创建鉴别器实例
discriminator = make_discriminator_model()

# %% 配置模型
cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)  # 定义二分类交叉熵损失函数

# 定义鉴别器损失函数
def discriminator_loss(real_output, fake_output):
    """
    计算鉴别器损失，包括对真实图像和生成图像的预测损失。
    """
    real_loss = cross_entropy(tf.ones_like(real_output), real_output)  # 真实图像的损失
    fake_loss = cross_entropy(tf.zeros_like(fake_output), fake_output)  # 生成图像的损失
    total_loss = real_loss + fake_loss  # 总损失
    return total_loss

# 定义生成器损失函数
def generator_loss(fake_output):
    """
    计算生成器损失，目标是最小化生成图像被判为真的概率。
    """
    return cross_entropy(tf.ones_like(fake_output), fake_output)  # 生成图像的损失

# 定义优化器
generator_optimizer = tf.keras.optimizers.Adam(1e-4)  # 生成器优化器
discriminator_optimizer = tf.keras.optimizers.Adam(1e-4)  # 鉴别器优化器

# 设置检查点
checkpoint_dir = './training_checkpoints'  # 检查点存储目录
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")  # 检查点前缀
checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer,
                                discriminator_optimizer=discriminator_optimizer,
                                generator=generator,
                                discriminator=discriminator)

# 如果存在检查点，则恢复最新的检查点
if os.path.exists(checkpoint_dir + "/checkpoint"):
    checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))

# 创建一组随机噪声向量用于生成示例图像
seed = tf.random.normal([num_examples_to_generate, noise_dim])

# 生成并保存图像
def generate_and_save_images(model, epoch, test_input):
    """
    使用生成器模型生成一批图像，并将它们保存到文件中。
    """
    predictions = model(test_input, training=False)  # 生成图像
    plt.figure(figsize=(4, 4))  # 设置绘图窗口大小
    for i in range(predictions.shape[0]):  # 遍历生成的图像
        plt.subplot(3, 3, i + 1)  # 创建子图
        plt.imshow(predictions[i, :, :, 0] * 127.5 + 127.5, cmap='gray')  # 显示图像
        plt.axis('off')  # 关闭坐标轴
    plt.savefig('image_at_epoch_{:04d}.png'.format(epoch))  # 保存图像到文件
    plt.show()  # 显示图像

# %% 定义训练步骤
@tf.function
def train_step(images):
    """
    定义单步训练过程，包括生成图像、计算损失和更新模型参数。
    """
    noise = tf.random.normal([BATCH_SIZE, noise_dim])  # 创建一批随机噪声向量

    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
        generated_images = generator(noise, training=True)  # 使用生成器生成图像

        real_output = discriminator(images, training=True)  # 使用鉴别器评估真实图像
        fake_output = discriminator(generated_images, training=True)  # 使用鉴别器评估生成图像

        gen_loss = generator_loss(fake_output)  # 计算生成器损失
        disc_loss = discriminator_loss(real_output, fake_output)  # 计算鉴别器损失

    # 计算梯度
    gradients_of_generator = gen_tape.gradient(gen_loss, generator.trainable_variables)
    gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.trainable_variables)

    # 应用梯度更新模型参数
    generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.trainable_variables))
    discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables))

    return gen_loss, disc_loss  # 返回损失值

# %% 定义训练过程
def train(dataset, epochs):
    """
    定义完整的训练过程，包括多个训练周期。
    """
    gen_losses = []  # 存储生成器损失
    disc_losses = []  # 存储鉴别器损失

    for epoch in range(epochs):
        start = time.time()  # 记录开始时间
        for image_batch in dataset:  # 遍历数据集中的每个批次
            gen_loss, disc_loss = train_step(image_batch)  # 执行单步训练
            gen_losses.append(gen_loss.numpy())  # 记录生成器损失
            disc_losses.append(disc_loss.numpy())  # 记录鉴别器损失

        # 清除输出，显示最新的生成图像
        display.clear_output(wait=True)
        generate_and_save_images(generator, epoch + 1, seed)  # 生成并保存图像

        # 每10个epoch保存一次模型
        if (epoch + 1) % 10 == 0:
            checkpoint.save(file_prefix=checkpoint_prefix)  # 保存检查点

        # 打印当前epoch的时间消耗
        print('Time for epoch {} is {} sec'.format(epoch + 1, time.time() - start))

    # 训练结束后绘制损失图
    plot_losses(gen_losses, disc_losses)

# 绘制损失图
def plot_losses(gen_losses, disc_losses):
    """
    绘制训练过程中生成器和鉴别器的损失曲线。
    """
    plt.figure(figsize=(12, 5))  # 设置绘图窗口大小
    plt.plot(gen_losses, label='Generator Loss')  # 绘制生成器损失曲线
    plt.plot(disc_losses, label='Discriminator Loss')  # 绘制鉴别器损失曲线
    plt.xlabel('Iteration')  # X轴标签
    plt.ylabel('Loss')  # Y轴标签
    plt.legend()  # 显示图例
    plt.title('Training Losses')  # 图表标题
    plt.savefig('loss_plot.png')  # 保存图表到文件
    plt.show()  # 显示图表

# %% 生成一个图像
def prediction():
    """
    使用训练好的生成器生成一张图像，并显示它。
    """
    test_input = tf.random.normal([1, noise_dim])  # 创建一个随机噪声向量
    predictions = generator(test_input, training=False)  # 使用生成器生成图像
    img_arr = predictions[0][:, :, 0]  # 提取生成的图像数据
    img_arr = img_arr * 127.5 + 127.5  # 将图像数据转换回原始像素值范围
    plt.axis('off')  # 关闭坐标轴
    plt.figure(figsize=(1, 1))  # 设置绘图窗口大小
    plt.imshow(img_arr, cmap='gray')  # 显示图像
    plt.show()  # 显示图像

    # 生成多张图像
    seed = tf.random.normal([num_examples_to_generate, noise_dim])  # 创建一组随机噪声向量
    generate_and_save_images(generator, 1, seed)  # 生成并保存图像

# %% 主程序入口
if __name__ == "__main__":
    print("welcome")  # 欢迎信息
    # 开始训练
    train(train_dataset, EPOCHS)
    # 生成图像
    prediction()